#include<iostream>
#include<cuda_runtime.h>


__device__ void preprocess_float(float* const data, int batch, int numData, int tidx,int tidy)
{
    if(tidy>numData) return;
    if(tidx>batch) return;    
    for(int i = tidy;i<numData;i+=blockDim.y)
    {
        unsigned int *data_temp = (unsigned int *)(&data[i + tidx*numData]);    
        *data_temp = (*data_temp >> 31 & 0x1)? ~(*data_temp): (*data_temp) | 0x80000000; 
    }
}

__device__ void Aeprocess_float(float* const data, int batch, int numData, int tidx,int tidy)
{
    for(int i = tidy;i<numData;i+=blockDim.y)
    {
        unsigned int* data_temp = (unsigned int *)(&data[i + tidx*numData]);
        *data_temp = (*data_temp >> 31 & 0x1)? (*data_temp) & 0x7fffffff: ~(*data_temp);
    }
}


__device__ void radixKernel(float* data0,float* data1,int tidx,int tidy,int batch,int numData){
    for(int bit=0;bit<32;bit++){
        unsigned int mask = 1 << bit;
        unsigned int cnt0 = 0,cnt1 = 0;   
        for(int i=tidy;i<numData;i+=blockDim.y){
            unsigned int *temp =(unsigned int *) &data0[i + tidx*numData];
            if(*temp&mask){
                data1[tidy+cnt1 + tidx*numData] = data0[i + tidx*numData];
                cnt1 += blockDim.y;
            }
            else{
                data0[tidy+cnt0+ tidx*numData] = data0[i+ tidx*numData];
                cnt0 += blockDim.y;
            }
        }    
        for(int j=0;j<cnt1;j+=blockDim.y){
            data0[j+cnt0+tidy+ tidx*numData] = data1[j+tidy+ tidx*numData];
        } 
    }
    return;
}

__device__ void mergeKernel(float* data0,float* data1,int* index,int tidx,int tidy,int batch,int numData){
    int numPerList = ceil((float)numData / blockDim.y);
    extern __shared__ int listIndexrecordValrecordTid[]; 
    int* listIndex = (int*)listIndexrecordValrecordTid; //记录线程上指针位置
    float* recordVal = (float*)listIndexrecordValrecordTid + batch*blockDim.y;  //得到应比较的数
    int * recordTid = (int*)listIndexrecordValrecordTid + 2*batch*blockDim.y;  // 记录当前线程
    int* recordSrcIndex = (int*)listIndexrecordValrecordTid + 3*batch*blockDim.y;  // 记录原index
    listIndex[tidy + tidx * blockDim.y] = 0;
    recordVal[tidy + tidx * blockDim.y] = 0;
    recordTid[tidy + tidx * blockDim.y] = tidy + tidx * blockDim.y;
    recordSrcIndex[tidy + tidx * blockDim.y] = 0;
    __syncthreads();

    for(int i=0;i<numData;i++){
        recordVal[tidy + tidx * blockDim.y] = 0;
        recordTid[tidy + tidx * blockDim.y] = tidy + tidx * blockDim.y;
        recordSrcIndex[tidy + tidx * blockDim.y] = 0; 
        if(listIndex[tidy + tidx * blockDim.y] < numPerList)
        {
            int src_index = tidy + tidx * numData + listIndex[tidy + tidx * blockDim.y]*blockDim.y;
            int batch_index = tidy + listIndex[tidy + tidx * blockDim.y]*blockDim.y;
            if(batch_index < numData)
            {
                recordVal[tidy + tidx * blockDim.y] = data0[src_index];
                recordSrcIndex[tidy + tidx * blockDim.y] = src_index;
            }
            else{
                unsigned int *temp = (unsigned int *)&recordVal[tidy + tidx * blockDim.y];
                *temp = 0xffffffff;
            }
        }else{
                unsigned int *temp = (unsigned int *)&recordVal[tidy + tidx * blockDim.y];
                *temp = 0xffffffff;
        }

        __syncthreads();
        int tidMax = blockDim.y >> 1;
        while (tidMax!=0)
        {
            if(tidy < tidMax)
            {
                unsigned int* temp1 = (unsigned int*)&recordVal[tidy + tidx * blockDim.y];
                unsigned int *temp2 = (unsigned int*)&recordVal[tidy + tidx * blockDim.y + tidMax];
                if(*temp2 < *temp1)
                {
                    recordVal[tidy + tidx * blockDim.y] = recordVal[tidy + tidx * blockDim.y + tidMax];
                    recordTid[tidy + tidx * blockDim.y] = recordTid[tidy + tidx * blockDim.y + tidMax];
                    recordSrcIndex[tidy + tidx * blockDim.y] = recordSrcIndex[tidy + tidx * blockDim.y + tidMax];
                }
            }
            tidMax = tidMax >> 1;
            __syncthreads();

        }
        if(tidy==0){
            listIndex[recordTid[tidx * blockDim.y]]++;
            data1[i + tidx * numData] = recordVal[tidx * blockDim.y];
            index[i + tidx * numData] = recordSrcIndex[tidx * blockDim.y]%numData;
        }
        __syncthreads();
        
    }
    return;
}

__global__ void radixSortGpu(float* src_data, float* dst_data, int* index,int batch, int dataLen){
    int tidx = threadIdx.x;
    int tidy = threadIdx.y;
    preprocess_float(src_data, batch,dataLen, tidx,tidy); 
    __syncthreads();
    radixKernel(src_data,dst_data,tidx,tidy,batch,dataLen);
    __syncthreads();
    mergeKernel(src_data,dst_data,index,tidx,tidy,batch,dataLen);
    __syncthreads();
    Aeprocess_float(dst_data, batch,dataLen, tidx,tidy);
    return;
}


int main(){
    using namespace std;
    int batch=4;
    int inputsLen = 40;
    int numBolcks=128;
    float* inputs;
    float* outputs;
    int* index;
    srand(100);    
    cudaMallocManaged(&inputs,sizeof(float)*inputsLen*batch);
    cudaMallocManaged(&outputs,sizeof(float)*inputsLen*batch);
    cudaMallocManaged(&index,sizeof(int)*inputsLen*batch);

    cout<<"input rand :"<<endl;
    for(int j=0;j<batch;j++){
        for(int i=0;i<inputsLen;i++){
            inputs[i+j*inputsLen] = (float)rand()/(float)RAND_MAX;
            index[i+j*inputsLen] = i;
            cout<<inputs[i+j*inputsLen]<<", ";
        }
        cout<<" "<<endl;
    }
    cout<<" "<<endl;

    const dim3 blockSize(batch,numBolcks);
    const dim3 girdSize(1,1);
    radixSortGpu<<<girdSize,blockSize,4*batch*numBolcks*sizeof(float)>>>(inputs,outputs,index,batch,inputsLen);
    cudaDeviceSynchronize();

    cout<<"output :"<<endl;
    for(int j=0;j<batch;j++){
        for(int i=0;i<inputsLen;i++){
            cout<<outputs[i+j*inputsLen]<<", ";
        }
        cout<<" "<<endl;
    }
    cout<<" "<<endl;   

    cout<<"index :"<<endl;
    for(int j=0;j<batch;j++){
        for(int i=0;i<inputsLen;i++){
            cout<<index[i+j*inputsLen]<<", ";
        }
        cout<<" "<<endl;
    }
    cout<<" "<<endl;  

    return 0;
}